"""
RNA seq data analysis
"""

import logging
import os
import time
import sys
import traceback
from argparse import ArgumentParser

# Add project path to system environment
PROJECT_DIR = os.path.abspath(os.path.join(os.getcwd(), "./"))
if PROJECT_DIR not in sys.path:
    sys.path.append(PROJECT_DIR)

from common import load_json, get_samples, get_cores, execute, save_list, generate_command



class RNASeq:
    def __init__(self, param, logger):
        self.config_data = param
        self.start_at = time.perf_counter()
        self.logger = logger
        self.dirs_name = ["1_clean", "2_alignment", "3_qualitative"]
        self.samples = get_samples(param["raw_reads_directory"], ".fq.gz", param["sample_name_rule"])
        self.cores = get_cores(self.config_data["core_usage"])
        print("Total number of sample is:", len(self.samples))


    def init_dirs(self):
        for dir_name in self.dirs_name:
            os.makedirs(f"{self.config_data['out_directory']}/{dir_name}", exist_ok=True)

    def write_log(self, log_item:dict):
        func = getattr(self.logger, log_item["type"])
        func(log_item["context"])

    def sequential(self, commands):
        for cmd in commands:
            exec_log = execute(cmd)
            self.write_log(exec_log)

    def clean(self):
        """
        Data quality control
        :return:
        """
        input_dir = self.config_data["raw_reads_directory"]
        output_dir = f"{self.config_data['out_directory']}/1_clean"
        commands = [
            generate_command({
                "commands": self.config_data["commands"],
                "program_name": self.config_data["steps"]["clean"],
                "input_dir": input_dir,
                "output_dir": output_dir,
                "options": self.config_data["options"]["clean"],
                "cores": self.cores,
                "sample_name": name,
                "mode": self.config_data['mode']
            },"clean")
            for name in self.samples]
        # Individual samples use the maximum number of processes
        if self.config_data['mode'] == 'HPC':
            # create dir for each samples
            for name in self.samples:
                os.makedirs(f"{output_dir}/{name}", exist_ok=True)

            save_list(f"{self.config_data['out_directory']}/clean_commands.txt", commands)
        else:
            self.sequential(commands)

    def alignment(self):
        """
        Alignment to the reference genome
        :return:
        """
        input_dir = f"{self.config_data['out_directory']}/1_clean"
        output_dir = f"{self.config_data['out_directory']}/2_alignment"
        commands = [
            generate_command({
                "commands": self.config_data["commands"],
                "program_name": self.config_data["steps"]["alignment"],
                "input_dir": f"{input_dir}/{name}" if self.config_data['mode'] == 'HPC' else input_dir,
                "output_dir": output_dir,
                "sample_name": name,
                "options": self.config_data["options"]["alignment"],
                "cores": self.cores,
                "hisat2_index_directory": self.config_data["hisat2_index_directory"]
            }, "alignment")
            for name in self.samples]
        if self.config_data['mode'] == 'HPC':
            save_list(f"{self.config_data['out_directory']}/alignment_commands.txt", commands)
        else:
            self.sequential(commands)


    def qualitative(self):
        """
        Quantitative analysis of the expression levels
        :return:
        """
        input_dir = f"{self.config_data['out_directory']}/2_alignment"
        output_dir = f"{self.config_data['out_directory']}/3_qualitative"

        if self.config_data["steps"]["qualitative"] == "stringtie":
            commands = [
                generate_command({
                    "commands": self.config_data["commands"],
                    "program_name": self.config_data["steps"]["qualitative"],
                    "input_dir": input_dir,
                    "output_dir": output_dir,
                    "options": self.config_data["options"]["qualitative"],
                    "sample_name": name,
                    "cores": self.cores,
                    "reference_annotation": self.config_data["reference_annotation"]
                }, "qualitative")
                for name in self.samples]
        else:
            commands = generate_command({
                "commands": self.config_data["commands"],
                "program_name": self.config_data["steps"]["qualitative"],
                "input_dir": input_dir,
                "output_dir": output_dir,
                "options": self.config_data["options"]["qualitative"],
                "cores": self.cores,
                "reference_annotation": self.config_data["reference_annotation"]
            }, "qualitative")

        if self.config_data['mode'] == 'HPC':
            save_list(f"{self.config_data['out_directory']}/qualitative_commands.txt", commands if type(commands) == list else [commands])
        else:
            self.sequential(commands if type(commands) == list else [commands])

    def runs(self):
        self.init_dirs()
        if "clean" in self.config_data["steps"]:
            self.clean()
            print(f"clean cost time is:{time.perf_counter() - self.start_at:.8f}s")
        if "alignment" in self.config_data["steps"]:
            self.alignment()
            print(f"alignment cost time is:{time.perf_counter() - self.start_at:.8f}s")
        if "qualitative" in self.config_data["steps"]:
            self.qualitative()
            print(f"qualitative cost time is:{time.perf_counter() - self.start_at:.8f}s")


if __name__ == "__main__":
    parser = ArgumentParser(description="")
    parser.add_argument(
        "-c", "--config", type=str, required=True, help=" specify config file "
    )
    args = parser.parse_args()
    config_data = load_json(args.config)

    log_name = int(time.time())
    log_file_path = f"{config_data['out_directory']}/{log_name}.log"
    logging.basicConfig(
        filename=log_file_path,
        filemode="a",
        level=logging.DEBUG,
        format="%(asctime)s - %(name)s - %(filename)s[line:%(lineno)d] - %(levelname)s - %(message)s",
    )
    logging.info("The program begins!")
    try:
        RNASeq(config_data, logging).runs()

    except Exception as e:
        exc_type, exc_value, exc_traceback_obj = sys.exc_info()
        print(exc_value)
        traceback.print_exception(
            exc_type,
            exc_value,
            exc_traceback_obj,
            file=open(log_file_path, "a", encoding="utf-8"),
        )

        sys.exit(exc_type)

    finally:
        logging.info("The program complete!")
